#!/bin/bash
# Copyright (c) 2022 Shenzhen Kaihong Digital Industry Development Co., Ltd. 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


cur_path=$(cd $(dirname $0); pwd)

NCNN_TARGET_ARCH="arm"

layer_list=(AbsVal ArgMax BatchNorm Bias BNLL Concat Convolution Crop Deconvolution Dropout Eltwise ELU Embed Exp Flatten InnerProduct Input Log LRN MemoryData MVN Pooling Power PReLU \
           Proposal Reduction ReLU Reshape ROIPooling Scale Sigmoid Slice Softmax Split SPP TanH Threshold Tile RNN LSTM BinaryOp UnaryOp ConvolutionDepthWise Padding Squeeze ExpandDims \
           Normalize Permute PriorBox DetectionOutput Interp DeconvolutionDepthWise ShuffleChannel InstanceNorm Clip Reorg YoloDetectionOutput Quantize Dequantize Yolov3DetectionOutput \
           PSROIPooling ROIAlign Packing Requantize Cast HardSigmoid SELU HardSwish Noop PixelShuffle DeepCopy Mish StatisticsPooling Swish Gemm GroupNorm LayerNorm Softplus GRU \
           MultiHeadAttention GELU Convolution1D Pooling1D ConvolutionDepthWise1D Convolution3D ConvolutionDepthWise3D Pooling3D)

head_files=(layer_declaration.h layer_registry.h layer_type_enum.h layer_shader_registry.h layer_shader_spv_data.h layer_shader_type_enum.h)

echo "layer_list=${layer_list}"

layer_declaration=""
layer_declaration_class=""
create_pipeline_content=""
destroy_pipeline_content=""
layer_registry=""
layer_type_enum=""

layer_shader_registry=""
layer_shader_spv_data=""
layer_shader_type_enum=""

__LAYER_TYPE_ENUM_INDEX=0

echo "Before get layer string"

for layer_name in ${layer_list[*]}; do

    layer_declaration_class=""    
    echo "layer_name=${layer_name}"
    if [[ ${layer_name} == "ArgMax" || ${layer_name} == "SPP" || ${layer_name} == "Tile" ]];then
        layer_registry=${layer_registry}"#if NCNN_STRING\n{\""${layer_name}"\", 0},\n#else\n{0},\n#endif\n"
    else
        #low_layer_name=`awk '{print tolower($layer_name)}'`
        low_layer_name=`echo ${layer_name,,}`
	echo "low_layer_name=${low_layer_name}"
        
	declare_str="#include \"layer/"${low_layer_name}".h\"\n"
	#echo "declare_str=${declare_str}"
        layer_declaration=${layer_declaration}${declare_str}
        
        class_declare_str="class "${layer_name}"_final : virtual public "${layer_name}
	#echo "class_declare_str=$class_declare_str"
        layer_declaration_class=${class_declare_str}

        create_pipeline_content="        { int ret = "${layer_name}"::create_pipeline(opt); if (ret) return ret; }\n"
        destroy_pipeline_content="        { int ret = "${layer_name}"::destroy_pipeline(opt); if (ret) return ret; }\n"
        
	arch_name=${cur_path}"/layer/"${NCNN_TARGET_ARCH}"/"${low_layer_name}"_"${NCNN_TARGET_ARCH}".cpp"
	echo "arch_name=${arch_name}"
        
	if [[ -f "${arch_name}" ]];then
	    echo "file ${arch_name} exist, need add declaration!"
	    arch_declare_str="#include \"layer/"${NCNN_TARGET_ARCH}"/"${low_layer_name}"_"${NCNN_TARGET_ARCH}".h\"\n"
	    layer_declaration=${layer_declaration}${arch_declare_str}

	    layer_declaration_class=${layer_declaration_class}", virtual public "${layer_name}"_"${NCNN_TARGET_ARCH}
	    create_pipeline_content=${create_pipeline_content}"        { int ret = "${layer_name}"_"${NCNN_TARGET_ARCH}"::create_pipeline(opt); if (ret) return ret; }\n"
	    destroy_pipeline_content="        { int ret = "${layer_name}"_"${NCNN_TARGET_ARCH}"::destroy_pipeline(opt); if (ret) return ret; }\n"${destroy_pipeline_content}
	else
	    echo "file ${arch_name} not exist!"

	fi
        
        layer_declaration=${layer_declaration}"namespace ncnn {\n"${layer_declaration_class}"\n{\n"
        layer_declaration=${layer_declaration}"public:\n"
        layer_declaration=${layer_declaration}"    virtual int create_pipeline(const Option& opt) {\n"${create_pipeline_content}"        return 0;\n    }\n"
        layer_declaration=${layer_declaration}"    virtual int destroy_pipeline(const Option& opt) {\n"${destroy_pipeline_content}"        return 0;\n    }\n"
        layer_declaration=${layer_declaration}"};\n"
        layer_declaration=${layer_declaration}"DEFINE_LAYER_CREATOR("${layer_name}"_final)\n} // namespace ncnn\n\n"

        layer_registry=${layer_registry}"#if NCNN_STRING\n{\""${layer_name}"\", "${layer_name}"_final_layer_creator},\n#else\n{"${layer_name}"_final_layer_creator},\n#endif\n"

    fi
    layer_type_enum=${layer_type_enum}${layer_name}" = "${__LAYER_TYPE_ENUM_INDEX}",\n"
    __LAYER_TYPE_ENUM_INDEX=$(( __LAYER_TYPE_ENUM_INDEX+1 ))
	
    echo -e "${layer_declaration}" > layer_declaration.txt
    echo -e "${layer_registry}" > layer_registry.txt
    echo -e "${layer_type_enum}" > layer_type_enum.txt
    echo "${__LAYER_TYPE_ENUM_INDEX}"
done

echo "all variables is get"

for file in ${head_files[*]}; do
    rm -rf $cur_path/${file}
    cp -rp $cur_path/${file}.in $cur_path/${file}
    
    line_num=`cat $cur_path/${file} | grep -n "@layer_declaration@" | awk -F ":" '{print $1}'`
    echo "file: ${file} find @layer_declaration@ line_num=${line_num}"
    if [[ ${line_num} != "" ]];then
        sed -i "${line_num}r layer_declaration.txt" $cur_path/${file}
        sed -i "${line_num}d" $cur_path/${file}
    fi
    
    line_num=`cat $cur_path/${file} | grep -n "@layer_registry@" | awk -F ":" '{print $1}'`
    echo "file: ${file} find @layer_registry@ line_num=${line_num}"
    if [[ ${line_num} != "" ]];then
        sed -i "${line_num}r layer_registry.txt" $cur_path/${file}
        sed -i "${line_num}d" $cur_path/${file}
    fi
   
    line_num=`cat $cur_path/${file} | grep -n "@layer_type_enum@" | awk -F ":" '{print $1}'`
    echo "file: ${file} find @layer_type_enum@ line_num=${line_num}"
    if [[ ${line_num} != "" ]];then
        sed -i "${line_num}r layer_type_enum.txt" $cur_path/${file}
        sed -i "${line_num}d" $cur_path/${file}
    fi

    line_num=`cat $cur_path/${file} | grep -n "@layer_shader_registry@" | awk -F ":" '{print $1}'`
    echo "file: ${file} find @layer_shader_registry@ line_num=${line_num}"
    if [[ ${line_num} != "" ]];then
        sed -i 's/@layer_shader_registry@/\/\/@layer_shader_registry@/g' $cur_path/${file}
        
    fi
    
    line_num=`cat $cur_path/${file} | grep -n "@layer_shader_spv_data@" | awk -F ":" '{print $1}'`
    echo "file: ${file} find @layer_shader_spv_data@ line_num=${line_num}"
    if [[ ${line_num} != "" ]];then
        sed -i 's/@layer_shader_spv_data@/\/\/@layer_shader_spv_data@/g' $cur_path/${file}

    fi

    line_num=`cat $cur_path/${file} | grep -n "@layer_shader_type_enum@" | awk -F ":" '{print $1}'`
    echo "file: ${file} find @layer_shader_type_enum@ line_num=${line_num}"
    if [[ ${line_num} != "" ]];then
        sed -i 's/@layer_shader_type_enum@/\/\/@layer_shader_type_enum@/g' $cur_path/${file}

    fi

    #sed -i 's/@layer_shader_registry@///@layer_shader_registry@/g' $cur_path/${file}
    #sed -i 's/@layer_shader_spv_data@///@layer_shader_spv_data@/g' $cur_path/${file}
    #sed -i 's/@layer_shader_type_enum@///@layer_shader_type_enum@/g' $cur_path/${file}
    
    echo "file:${file} is replaced"
done


